#  Copyright (c) 2010
#  The Regents of the University of Michigan
#  All Rights Reserved

#  Permission is granted to use, copy, create derivative works, and
#  redistribute this software and such derivative works for any purpose,
#  so long as the name of the University of Michigan is not used in
#  any advertising or publicity pertaining to the use or distribution
#  of this software without specific, written prior authorization. If
#  the above copyright notice or any other identification of the
#  University of Michigan is included in any copy of any portion of
#  this software, then the disclaimer below must also be included.

#  This software is provided as is, without representation or warranty
#  of any kind either express or implied, including without limitation
#  the implied warranties of merchantability, fitness for a particular
#  purpose, or noninfringement.  The Regents of the University of
#  Michigan shall not be liable for any damages, including special,
#  indirect, incidental, or consequential damages, with respect to any
#  claim arising out of or in connection with the use of the software,
#  even if it has been or is hereafter advised of the possibility of
#  such damages.
import matplotlib
matplotlib.use('cairo.png')
from matplotlib import pyplot

import collections
import cPickle as pickle
import glob
import numpy
import os
import sys

from util import *

class TorrentMetadataDump(object):
  name = 'metadata'

  chart_title = 'Torrents Per Crawl with .torrent Downloaded'
  @staticmethod
  def get_num_torrents(dir):
    return sum(1 for x in os.listdir(dir) if '.torrent' in x)

  @staticmethod
  def get_torrents(dir):
    return frozenset(x for x in os.listdir(dir) if '.torrent' in x)

  @staticmethod
  def get_unique_torrents(dir):
    known_torrents = collections.defaultdict(int)
    for dir in dirs:
      for torrent in strategy.get_torrents(dir):
        known_torrents[torrent] += 1

    unique_torrents = []
    for dir in dirs:
      unique_torrents.append(sum(1 for torrent in strategy.get_torrents(dir)
                                 if known_torrents[torrent] == 1))

    return unique_torrents


class DatabaseStrategy(object):
  def __init__(self):
    self.cached_count = {}
    self.cached_torrents = {}

  def get_num_torrents(self, dir):
    if self.cache:
      return len(self.get_torrents(dir))
    elif self.cached_count.get(dir) is None:
      s = get_db_suffix(dir)
      conn = get_db_conn()
      c = conn.cursor()
      c.execute(self.count_query % {'suffix': s})
      self.cached_count[dir] = c.fetchone()[0]
    return self.cached_count[dir]


  def get_torrents(self, dir):
    if self.cached_torrents.get(dir) is None:
      s = get_db_suffix(dir)
      conn = get_db_conn()
      c = conn.cursor()
      c.execute(self.query % {'suffix': s})
      self.cached_torrents[dir] = frozenset(c.fetchall())
    return self.cached_torrents[dir]


class RelatedContentAndPeerLists(DatabaseStrategy):
  name = 'rel_and_peers'

  query = """SELECT * FROM
             (SELECT D.hash
              FROM torrent_descs_%(suffix)s D, peer_lists_%(suffix)s P
              WHERE digest(D.hash, 'sha1') = P.dhtkey
              GROUP BY D.hash) s;"""

  chart_title = 'Torrents per crawl with descriptions and peer lists'
  

class RelatedContent(DatabaseStrategy):
  name = 'rel'

  cache = False

  count_query = '''SELECT COUNT(*) FROM (SELECT hash
                            FROM torrent_descs_%(suffix)s
                            WHERE hash IS NOT NULL
                            GROUP BY hash) s;'''

  query = '''SELECT * FROM (SELECT hash
                            FROM torrent_descs_%(suffix)s
                            WHERE hash IS NOT NULL
                            GROUP BY hash) s;'''

  chart_title = 'Torrents per crawl with descriptions'
  

class PeerLists(DatabaseStrategy):
  name = 'peer_lists'

  cache = False

  count_query = '''SELECT COUNT(*)
             FROM (SELECT dhtkey FROM peer_lists_%(suffix)s
                   GROUP BY dhtkey) s;'''
  query = '''SELECT *
             FROM (SELECT dhtkey FROM peer_lists_%(suffix)s
                   GROUP BY dhtkey) s;'''

  chart_title = 'Torrents Per Crawl with Peer Lists'


def calc_unique_torrents(dirs, strategy):
  unique_torrents = set()
  not_unique_torrents = set()

  for dir in dirs:
    for torrent in strategy.get_torrents(dir):
      if torrent in unique_torrents:
        unique_torrents.remove(torrent)
        not_unique_torrents.add(torrent)
      elif torrent not in not_unique_torrents:
        unique_torrents.add(torrent)

  unique_torrents_counts = []
  for dir in dirs:
    dircount = sum(1 for torrent in strategy.get_torrents(dir)
                   if torrent in unique_torrents)
    unique_torrents_counts.append(dircount)

  return unique_torrents_counts


def make_plot(dirs, strategy):
  pyplot.clf()
  ind = numpy.arange(len(dirs))

  try:
    total_torrents = pickle.load(open('torrents_total_%s.dat' % strategy.name, 'r'))
  except IOError:
    total_torrents = [strategy.get_num_torrents(x) for x in dirs]
    pickle.dump(total_torrents,
                open('torrents_total_%s.dat' % strategy.name, 'w'),
                pickle.HIGHEST_PROTOCOL)

  pt1 = pyplot.bar(ind, total_torrents)
  del total_torrents

  try:
    new_torrents = pickle.load(open('torrents_new_%s.dat' % strategy.name, 'r'))
  except IOError:
    # New torrents.
    known_torrents = set()
    last_size = 0
    new_torrents = []
    for dir in dirs:
      known_torrents.update(strategy.get_torrents(dir))
      new_size = len(known_torrents)
      new_torrents.append(new_size - last_size)
      last_size = new_size

    del known_torrents
    pickle.dump(new_torrents, open('torrents_new_%s.dat' % strategy.name, 'w'),
                pickle.HIGHEST_PROTOCOL)
    
  pt2 = pyplot.bar(ind, new_torrents, color='r')
  del new_torrents

  try:
    unique_torrents = pickle.load(open('torrents_uniq_%s.dat' % strategy.name, 'r'))
  except IOError:
    unique_torrents = calc_unique_torrents(dirs, strategy)
    pickle.dump(unique_torrents,
                open('torrents_uniq_%s.dat' % strategy.name, 'w'),
                pickle.HIGHEST_PROTOCOL)

  pt3 = pyplot.bar(ind, unique_torrents, color='g')
  del unique_torrents

  pyplot.xlabel('Crawl')
  pyplot.title(strategy.chart_title)
  pyplot.ylabel('Number of torrents')

  pyplot.figlegend((pt1[0], pt2[0], pt3[0]),
                   ('Already known',  'Discovered', 'Discovered only once'),
                   loc='upper right')
  pyplot.savefig('torrents_%s.png' % strategy.name)


def main(argv):
  strategies = dict((x.name, x) for x in (TorrentMetadataDump, PeerLists,
                    RelatedContent, RelatedContentAndPeerLists))
  dirs = glob.glob(argv[0])
  dirs.sort(key=get_date)
  if argv[1] == 'db':
    for strategy in strategies.values():
      if strategy.name != 'metadata':
        make_plot(dirs, strategy())
  else:
    make_plot(dirs, strategies[argv[1]]())

if __name__ == '__main__':
  main(sys.argv[1:])
